from celery import Celery
from app import create_app
from kombu import Exchange, Queue

def create_celery():
    app = create_app()
    
    # 创建优先级队列
    task_queues = [
        Queue('high', Exchange('high'), routing_key='high',
              queue_arguments={'x-max-priority': 10}),
        Queue('default', Exchange('default'), routing_key='default',
              queue_arguments={'x-max-priority': 10}),
        Queue('low', Exchange('low'), routing_key='low',
              queue_arguments={'x-max-priority': 10}),
    ]
    
    celery = Celery(
        app.import_name,
        broker=app.config['CELERY']['broker_url'],
        backend=app.config['CELERY']['result_backend']
    )
    
    celery.conf.update(app.config['CELERY'])
    celery.conf.task_queues = task_queues
    celery.conf.task_default_queue = 'default'
    celery.conf.task_default_exchange = 'default'
    celery.conf.task_default_routing_key = 'default'
    
    # 任务路由
    celery.conf.task_routes = {
        'app.tasks.document_tasks.process_document': {'queue': 'high'},
        'app.tasks.document_tasks.generate_embeddings': {'queue': 'default'},
    }

    class ContextTask(celery.Task):
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return self.run(*args, **kwargs)

    celery.Task = ContextTask
    return celery

celery = create_celery() 